using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

namespace GenerativeAgents.Core
{
    /// <summary>
    /// 检索引擎
    /// 负责从记忆系统中高效检索相关信息以支持智能体决策
    /// </summary>
    public class RetrievalEngine : MonoBehaviour
    {
        [Header("检索配置")]
        [SerializeField] private float recencyWeight = 0.3f;
        [SerializeField] private float importanceWeight = 0.3f;
        [SerializeField] private float relevanceWeight = 0.4f;
        [SerializeField] private int defaultRetrievalCount = 10;
        [SerializeField] private bool enableRetrievalLogging = false;
        
        [Header("高级检索设置")]
        [SerializeField] private float semanticThreshold = 0.5f;
        [SerializeField] private float temporalDecayRate = 0.01f;
        [SerializeField] private bool useContextualBoost = true;
        [SerializeField] private float contextualBoostFactor = 0.2f;
        
        // 组件引用
        private MemorySystem memorySystem;
        private IEmbeddingEngine embeddingEngine;
        private QueryProcessor queryProcessor;
        private ScoreCalculator scoreCalculator;
        
        // 检索统计
        private int totalQueries = 0;
        private int totalRetrievedMemories = 0;
        private Dictionary<QueryContext, int> queryTypeStats;
        private List<string> recentQueries;
        
        // 缓存
        private Dictionary<string, List<MemoryNode>> queryCache;
        private int maxCacheSize = 100;
        
        private void Awake()
        {
            InitializeRetrievalEngine();
        }
        
        private void InitializeRetrievalEngine()
        {
            // 获取组件引用
            memorySystem = GetComponent<MemorySystem>();
            embeddingEngine = GetComponent<IEmbeddingEngine>();
            
            // 创建子组件
            queryProcessor = new QueryProcessor();
            scoreCalculator = new ScoreCalculator();
            
            // 初始化数据结构
            queryTypeStats = new Dictionary<QueryContext, int>();
            foreach (QueryContext context in Enum.GetValues(typeof(QueryContext)))
            {
                queryTypeStats[context] = 0;
            }
            
            recentQueries = new List<string>();
            queryCache = new Dictionary<string, List<MemoryNode>>();
            
            if (enableRetrievalLogging)
            {
                Debug.Log("[RetrievalEngine] Retrieval engine initialized");
            }
        }
        
        /// <summary>
        /// 检索相关记忆（基于文本查询）
        /// </summary>
        public List<MemoryNode> RetrieveRelevant(string query, QueryContext context = QueryContext.General, int maxResults = -1)
        {
            if (maxResults <= 0) maxResults = defaultRetrievalCount;
            
            try
            {
                // 处理查询
                var processedQuery = queryProcessor.ProcessQuery(query, context);
                
                // 检查缓存
                var cacheKey = GenerateCacheKey(processedQuery);
                if (queryCache.TryGetValue(cacheKey, out var cachedResults))
                {
                    if (enableRetrievalLogging)
                    {
                        Debug.Log($"[RetrievalEngine] Retrieved {cachedResults.Count} memories from cache for query: {query}");
                    }
                    return cachedResults.Take(maxResults).ToList();
                }
                
                // 执行检索
                var retrievedMemories = PerformRetrieval(processedQuery, maxResults);
                
                // 缓存结果
                CacheResults(cacheKey, retrievedMemories);
                
                // 更新统计
                UpdateRetrievalStats(query, context, retrievedMemories.Count);
                
                if (enableRetrievalLogging)
                {
                    Debug.Log($"[RetrievalEngine] Retrieved {retrievedMemories.Count} memories for query: {query}");
                }
                
                return retrievedMemories;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RetrievalEngine] Error retrieving memories for query '{query}': {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 检索用于计划的记忆
        /// </summary>
        public List<MemoryNode> RetrieveForPlanning(string goal, Vector3 currentLocation, string agentName)
        {
            try
            {
                var planningQuery = MemoryQuery.CreatePlanningQuery(goal, currentLocation);
                planningQuery.agentName = agentName;
                
                // 增加计划相关的类型过滤
                planningQuery.typeFilter.Add(MemoryType.Plan);
                planningQuery.typeFilter.Add(MemoryType.Event);
                planningQuery.typeFilter.Add(MemoryType.Thought);
                
                // 调整权重，更重视重要性和相关性
                planningQuery.importanceWeight = 0.4f;
                planningQuery.relevanceWeight = 0.4f;
                planningQuery.recencyWeight = 0.2f;
                
                var retrievedMemories = memorySystem.RetrieveMemories(planningQuery);
                
                // 添加上下文相关的记忆
                if (useContextualBoost)
                {
                    var contextualMemories = RetrieveContextualMemories(goal, currentLocation, agentName);
                    retrievedMemories = MergeAndRankMemories(retrievedMemories, contextualMemories, planningQuery.maxResults);
                }
                
                UpdateRetrievalStats(goal, QueryContext.Planning, retrievedMemories.Count);
                
                return retrievedMemories;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RetrievalEngine] Error retrieving memories for planning '{goal}': {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 检索用于对话的记忆
        /// </summary>
        public List<MemoryNode> RetrieveForConversation(string topic, string conversationPartner, string agentName)
        {
            try
            {
                var conversationQuery = MemoryQuery.CreateConversationQuery(topic, conversationPartner);
                conversationQuery.agentName = agentName;
                
                // 添加对话相关的记忆检索
                var topicMemories = memorySystem.RetrieveMemories(conversationQuery);
                
                // 检索与对话伙伴相关的记忆
                var partnerQuery = new MemoryQuery
                {
                    queryText = conversationPartner,
                    maxResults = 5,
                    context = QueryContext.Social
                };
                partnerQuery.participantFilter.Add(conversationPartner);
                
                var partnerMemories = memorySystem.RetrieveMemories(partnerQuery);
                
                // 合并和排序记忆
                var allMemories = MergeAndRankMemories(topicMemories, partnerMemories, conversationQuery.maxResults);
                
                UpdateRetrievalStats(topic, QueryContext.Conversation, allMemories.Count);
                
                return allMemories;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RetrievalEngine] Error retrieving memories for conversation about '{topic}': {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 检索用于反思的记忆
        /// </summary>
        public List<MemoryNode> RetrieveForReflection(List<string> focusPoints, DateTime timeStart, string agentName)
        {
            try
            {
                var allReflectionMemories = new List<MemoryNode>();
                
                foreach (var focusPoint in focusPoints)
                {
                    var reflectionQuery = MemoryQuery.CreateReflectionQuery(focusPoint, timeStart);
                    reflectionQuery.agentName = agentName;
                    reflectionQuery.maxResults = 15; // 每个焦点点获取更多记忆
                    
                    var focusMemories = memorySystem.RetrieveMemories(reflectionQuery);
                    allReflectionMemories.AddRange(focusMemories);
                }
                
                // 去重并按综合分数排序
                var uniqueMemories = allReflectionMemories
                    .GroupBy(m => m.nodeId)
                    .Select(g => g.First())
                    .ToList();
                
                var rankedMemories = RankMemoriesForReflection(uniqueMemories, focusPoints);
                
                UpdateRetrievalStats(string.Join(", ", focusPoints), QueryContext.Reflection, rankedMemories.Count);
                
                return rankedMemories.Take(20).ToList(); // 限制反思记忆数量
            }
            catch (Exception ex)
            {
                Debug.LogError($"[RetrievalEngine] Error retrieving memories for reflection: {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 执行基础检索
        /// </summary>
        private List<MemoryNode> PerformRetrieval(ProcessedQuery processedQuery, int maxResults)
        {
            if (memorySystem == null)
            {
                Debug.LogWarning("[RetrievalEngine] Memory system not available");
                return new List<MemoryNode>();
            }
            
            // 创建内存查询对象
            var memoryQuery = CreateMemoryQuery(processedQuery, maxResults);
            
            // 执行检索
            var memories = memorySystem.RetrieveMemories(memoryQuery);
            
            // 应用上下文增强
            if (useContextualBoost && processedQuery.context != QueryContext.General)
            {
                memories = ApplyContextualBoost(memories, processedQuery);
            }
            
            return memories;
        }
        
        /// <summary>
        /// 创建内存查询对象
        /// </summary>
        private MemoryQuery CreateMemoryQuery(ProcessedQuery processedQuery, int maxResults)
        {
            var memoryQuery = new MemoryQuery
            {
                queryText = processedQuery.originalQuery,
                queryEmbedding = processedQuery.embedding,
                keywords = processedQuery.keywords,
                maxResults = maxResults,
                recencyWeight = recencyWeight,
                importanceWeight = importanceWeight,
                relevanceWeight = relevanceWeight,
                context = processedQuery.context
            };
            
            // 根据查询上下文设置过滤器
            ApplyContextFilters(memoryQuery, processedQuery.context);
            
            return memoryQuery;
        }
        
        /// <summary>
        /// 应用上下文过滤器
        /// </summary>
        private void ApplyContextFilters(MemoryQuery memoryQuery, QueryContext context)
        {
            switch (context)
            {
                case QueryContext.Conversation:
                    memoryQuery.typeFilter.Add(MemoryType.Chat);
                    memoryQuery.typeFilter.Add(MemoryType.Event);
                    memoryQuery.typeFilter.Add(MemoryType.Thought);
                    break;
                    
                case QueryContext.Planning:
                    memoryQuery.typeFilter.Add(MemoryType.Plan);
                    memoryQuery.typeFilter.Add(MemoryType.Event);
                    memoryQuery.typeFilter.Add(MemoryType.Knowledge);
                    break;
                    
                case QueryContext.Reflection:
                    memoryQuery.typeFilter.Add(MemoryType.Event);
                    memoryQuery.typeFilter.Add(MemoryType.Thought);
                    memoryQuery.typeFilter.Add(MemoryType.Reflection);
                    // 增加重要性权重
                    memoryQuery.importanceWeight = 0.5f;
                    memoryQuery.recencyWeight = 0.3f;
                    memoryQuery.relevanceWeight = 0.2f;
                    break;
                    
                case QueryContext.Social:
                    memoryQuery.typeFilter.Add(MemoryType.Chat);
                    memoryQuery.typeFilter.Add(MemoryType.Event);
                    memoryQuery.typeFilter.Add(MemoryType.Observation);
                    break;
            }
        }
        
        /// <summary>
        /// 应用上下文增强
        /// </summary>
        private List<MemoryNode> ApplyContextualBoost(List<MemoryNode> memories, ProcessedQuery processedQuery)
        {
            var boostedMemories = new List<(MemoryNode memory, float boostedScore)>();
            
            foreach (var memory in memories)
            {
                var baseScore = scoreCalculator.CalculateCompositeScore(memory, processedQuery, recencyWeight, importanceWeight, relevanceWeight);
                var contextBoost = CalculateContextualBoost(memory, processedQuery);
                var boostedScore = baseScore + (contextBoost * contextualBoostFactor);
                
                boostedMemories.Add((memory, boostedScore));
            }
            
            return boostedMemories
                .OrderByDescending(x => x.boostedScore)
                .Select(x => x.memory)
                .ToList();
        }
        
        /// <summary>
        /// 计算上下文增强分数
        /// </summary>
        private float CalculateContextualBoost(MemoryNode memory, ProcessedQuery processedQuery)
        {
            float boost = 0f;
            
            switch (processedQuery.context)
            {
                case QueryContext.Conversation:
                    if (memory.type == MemoryType.Chat) boost += 0.3f;
                    if (memory.participants?.Count > 1) boost += 0.2f;
                    break;
                    
                case QueryContext.Planning:
                    if (memory.type == MemoryType.Plan) boost += 0.4f;
                    if (memory.poignancy >= 7) boost += 0.2f;
                    break;
                    
                case QueryContext.Reflection:
                    if (memory.type == MemoryType.Thought || memory.type == MemoryType.Reflection) boost += 0.3f;
                    if (memory.accessCount > 3) boost += 0.2f;
                    break;
                    
                case QueryContext.Social:
                    if (memory.participants?.Count > 0) boost += 0.2f;
                    if (memory.type == MemoryType.Event || memory.type == MemoryType.Chat) boost += 0.1f;
                    break;
            }
            
            return boost;
        }
        
        /// <summary>
        /// 检索上下文相关记忆
        /// </summary>
        private List<MemoryNode> RetrieveContextualMemories(string goal, Vector3 location, string agentName)
        {
            var contextualMemories = new List<MemoryNode>();
            
            // 检索位置相关的记忆
            var locationQuery = new MemoryQuery
            {
                queryText = "location activities",
                locationCenter = location,
                locationRadius = 10f,
                maxResults = 5
            };
            
            var locationMemories = memorySystem.RetrieveMemories(locationQuery);
            contextualMemories.AddRange(locationMemories);
            
            // 检索最近的高重要性记忆
            var recentImportantMemories = memorySystem.AssociativeMemory.GetHighImportanceMemories(6, 5);
            contextualMemories.AddRange(recentImportantMemories);
            
            return contextualMemories.Distinct().ToList();
        }
        
        /// <summary>
        /// 合并和排序记忆
        /// </summary>
        private List<MemoryNode> MergeAndRankMemories(List<MemoryNode> memories1, List<MemoryNode> memories2, int maxResults)
        {
            var allMemories = memories1.Concat(memories2)
                .GroupBy(m => m.nodeId)
                .Select(g => g.First())
                .ToList();
            
            // 按综合重要性排序
            return allMemories
                .OrderByDescending(m => m.GetCompositeImportance())
                .ThenByDescending(m => m.lastAccessTime)
                .Take(maxResults)
                .ToList();
        }
        
        /// <summary>
        /// 为反思排序记忆
        /// </summary>
        private List<MemoryNode> RankMemoriesForReflection(List<MemoryNode> memories, List<string> focusPoints)
        {
            var rankedMemories = new List<(MemoryNode memory, float score)>();
            
            foreach (var memory in memories)
            {
                var score = memory.GetCompositeImportance();
                
                // 基于焦点点增加分数
                foreach (var focusPoint in focusPoints)
                {
                    if (memory.content.Contains(focusPoint, StringComparison.OrdinalIgnoreCase) ||
                        memory.keywords?.Any(k => k.Contains(focusPoint, StringComparison.OrdinalIgnoreCase)) == true)
                    {
                        score += 0.3f;
                    }
                }
                
                // 反思类型的记忆获得额外分数
                if (memory.type == MemoryType.Thought || memory.type == MemoryType.Reflection)
                {
                    score += 0.2f;
                }
                
                rankedMemories.Add((memory, score));
            }
            
            return rankedMemories
                .OrderByDescending(x => x.score)
                .Select(x => x.memory)
                .ToList();
        }
        
        /// <summary>
        /// 生成缓存键
        /// </summary>
        private string GenerateCacheKey(ProcessedQuery processedQuery)
        {
            return $"{processedQuery.originalQuery}_{processedQuery.context}_{string.Join(",", processedQuery.keywords)}";
        }
        
        /// <summary>
        /// 缓存结果
        /// </summary>
        private void CacheResults(string cacheKey, List<MemoryNode> results)
        {
            if (queryCache.Count >= maxCacheSize)
            {
                // 移除最旧的缓存项
                var oldestKey = queryCache.Keys.First();
                queryCache.Remove(oldestKey);
            }
            
            queryCache[cacheKey] = new List<MemoryNode>(results);
        }
        
        /// <summary>
        /// 更新检索统计
        /// </summary>
        private void UpdateRetrievalStats(string query, QueryContext context, int resultCount)
        {
            totalQueries++;
            totalRetrievedMemories += resultCount;
            queryTypeStats[context]++;
            
            // 记录最近查询
            recentQueries.Add(query);
            if (recentQueries.Count > 50)
            {
                recentQueries.RemoveAt(0);
            }
        }
        
        /// <summary>
        /// 清理缓存
        /// </summary>
        public void ClearCache()
        {
            queryCache.Clear();
        }
        
        /// <summary>
        /// 获取检索统计
        /// </summary>
        public string GetRetrievalStats()
        {
            var avgResults = totalQueries > 0 ? (float)totalRetrievedMemories / totalQueries : 0f;
            var stats = $"Total Queries: {totalQueries}, Avg Results: {avgResults:F1}\n";
            
            stats += "Query Types:\n";
            foreach (var kvp in queryTypeStats)
            {
                if (kvp.Value > 0)
                {
                    stats += $"  {kvp.Key}: {kvp.Value}\n";
                }
            }
            
            stats += $"Cache Size: {queryCache.Count}/{maxCacheSize}";
            
            return stats;
        }
        
        /// <summary>
        /// 设置检索参数
        /// </summary>
        public void SetRetrievalParameters(float[] weights, int defaultCount, float threshold)
        {
            if (weights.Length >= 3)
            {
                recencyWeight = weights[0];
                importanceWeight = weights[1];
                relevanceWeight = weights[2];
            }
            
            defaultRetrievalCount = defaultCount;
            semanticThreshold = threshold;
        }
        
        /// <summary>
        /// 重置统计信息
        /// </summary>
        public void ResetStats()
        {
            totalQueries = 0;
            totalRetrievedMemories = 0;
            
            foreach (var key in queryTypeStats.Keys.ToList())
            {
                queryTypeStats[key] = 0;
            }
            
            recentQueries.Clear();
        }
    }
}